package com.jianjun.common.util;


import com.jfinal.kit.LogKit;
import com.jfinal.kit.StrKit;
import com.jfinal.plugin.activerecord.ActiveRecordException;
import com.jfinal.plugin.activerecord.Config;
import com.jfinal.plugin.activerecord.DbKit;
import com.jfinal.plugin.activerecord.IAtom;
import com.jfinal.plugin.activerecord.NestedTransactionHelpException;

import javax.sql.DataSource;
import javax.sql.XAConnection;
import javax.sql.XADataSource;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * XA 事务
 */
public class DbXA {

    public static boolean tx(IAtom atom, String... configNames) {
        return tx(Connection.TRANSACTION_SERIALIZABLE, atom, configNames);
    }

    public static boolean tx(int transactionLevel, IAtom atom, String... configNames) {
        List<DbXAUnit> list = createDbXAUnitList(configNames);
        try {
            // 初始化XA链接
            init(transactionLevel, list);
            // 开始XA事务
            start(list);
            // 执行SQL操作
            boolean result = atom.run();
            if (!result) {
                // 回滚事务
                rollback(list);
                return false;
            }
            // 结束事务的SQL操作阶段
            end(list);
            // 准备提交事务
            if (prepare(list)) {
                // 提交事务
                commit(list);
                return true;
            } else {
                // 回滚事务
                rollback(list);
                return false;
            }
        } catch (NestedTransactionHelpException e) {
            // 回滚事务
            rollback(list);
            LogKit.logNothing(e);
            return false;
        } catch (Throwable t) {
            // 回滚事务
            rollback(list);
            throw t instanceof RuntimeException ? (RuntimeException) t : new ActiveRecordException(t);
        } finally {
            // 关闭连接
            close(list);
        }
    }

    private static void init(int transactionLevel, List<DbXAUnit> list) throws SQLException {
        DbXidGroup dbXidGroup = new DbXidGroup(transactionLevel);
        for (DbXAUnit unit : list) {
            unit.init(dbXidGroup.create(unit.configName));
        }
    }

    private static void start(List<DbXAUnit> list) throws XAException {
        for (DbXAUnit unit : list) {
            unit.xaResource.start(unit.xid, XAResource.TMJOIN);
        }
    }

    private static void end(List<DbXAUnit> list) throws XAException {
        for (DbXAUnit unit : list) {
            unit.xaResource.end(unit.xid, XAResource.TMSUCCESS);
        }
    }

    private static void commit(List<DbXAUnit> list) throws XAException {
        for (DbXAUnit unit : list) {
            unit.xaResource.commit(unit.xid, false);
            //unit.config.executeCallbackAfterTxCommit();
        }
    }

    private static void close(List<DbXAUnit> list) {
        for (DbXAUnit unit : list) {
            unit.close();
        }
    }

    private static void rollback(List<DbXAUnit> list) {
        for (DbXAUnit unit : list) {
            unit.rollback();
        }
    }

    private static boolean prepare(List<DbXAUnit> list) throws XAException {
        for (DbXAUnit unit : list) {
            int prepare = unit.xaResource.prepare(unit.xid);
            if (prepare != XAResource.XA_OK && prepare != XAResource.XA_RDONLY) {
                return false;
            }
        }
        return true;
    }

    private static List<DbXAUnit> createDbXAUnitList(String[] configNames) {
        List<DbXAUnit> list = new ArrayList<>(configNames.length);
        for (String configName : configNames) {
            list.add(createDbXAUnit(configName));
        }
        return list;
    }

    private static DbXAUnit createDbXAUnit(String configName) {
        Config config = DbKit.getConfig(configName);
        DataSource dataSource = config.getDataSource();
        if (!(dataSource instanceof XADataSource)) {
            throw new RuntimeException(configName + " : dataSource is not XADataSource[数据库连接池不支持XA模式]");
        }
        return new DbXAUnit(config, (XADataSource) dataSource);
    }

    public static class DbXAUnit {
        private final Config config;
        private final String configName;
        private final XADataSource xaDataSource;
        private XAConnection xaConnection;
        private XAResource xaResource;
        private Xid xid;

        public DbXAUnit(Config config, XADataSource xaDataSource) {
            this.config = config;
            this.configName = config.getName();
            this.xaDataSource = xaDataSource;
        }

        public void init(Xid xid) throws SQLException {
            this.xid = xid;
            xaConnection = xaDataSource.getXAConnection();
            xaResource = xaConnection.getXAResource();
            config.setThreadLocalConnection(xaConnection.getConnection());
        }

        public void close() {
            if (this.xaConnection != null) {
                try {
                    this.xaConnection.close();
                } catch (Throwable t) {
                    LogKit.error("关闭xa连接异常:" + configName + " : " + xid.toString() + " ，" + t.getMessage(), t);
                } finally {
                    config.removeThreadLocalConnection();
                }
            }
        }

        public void rollback() {
            if (xaResource != null) {
                try {
                    xaResource.rollback(xid);
                } catch (Throwable t) {
                    LogKit.error("回滚xa连接异常:" + configName + " : " + xid.toString() + " ，" + t.getMessage(), t);
                }
            }
        }

        @Override
        public String toString() {
            return "DbXAGroup{" +
                    "configName='" + configName + '\'' +
                    ", xaDataSource=" + xaDataSource +
                    ", xaConnection=" + xaConnection +
                    ", xaResource=" + xaResource +
                    ", xid=" + xid +
                    '}';
        }
    }

    public static class DbXidGroup {
        private final String groupID = StrKit.getRandomUUID();
        private final byte[] globalTransactionId = groupID.getBytes();
        private final int formatId;

        public DbXidGroup(int formatId) {
            this.formatId = formatId;
        }

        public Xid create(String configName) {
            byte[] branchQualifier = configName.getBytes();
            return new Xid() {
                @Override
                public String toString() {
                    return "Xid{" +
                            "formatId=" + formatId +
                            ", groupID='" + groupID + '\'' +
                            ", configName='" + configName + '\'' +
                            '}';
                }

                @Override
                public int getFormatId() {
                    return formatId;
                }

                @Override
                public byte[] getGlobalTransactionId() {
                    return globalTransactionId;
                }

                @Override
                public byte[] getBranchQualifier() {
                    return branchQualifier;
                }
            };
        }
    }

}